#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Author: kerlomz <kerlomz@gmail.com>
import io
import cv2
import time
import PIL.Image as PilImage
import numpy as np
import tensorflow as tf
from enum import Enum, unique
from distutils.version import StrictVersion


@unique
class TargetColor(Enum):
    Black = 0
    Red = 1
    Blue = 2
    Yellow = 3
    Green = 4
    White = 7


color_map = {
    'black': TargetColor.Black,
    'red': TargetColor.Red,
    'blue': TargetColor.Blue,
    'yellow': TargetColor.Yellow,
    'green': TargetColor.Green,
    'white': TargetColor.White
}


class ColorExtract:

    def __init__(self):
        self.model_raw_v1_14 = b'\nB\n\x05Const\x12\x05Const*%\n\x05value\x12\x1cB\x1a\x08\x03\x12\x08\x12\x02\x08\x01\x12\x02\x08\x03"\x0c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nD\n\x07Const_1\x12\x05Const*%\n\x05value\x12\x1cB\x1a\x08\x03\x12\x08\x12\x02\x08\x01\x12\x02\x08\x03"\x0c\x00\x00\x00\x00\x00\x00\x00\x00\xff\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nD\n\x07Const_2\x12\x05Const*%\n\x05value\x12\x1cB\x1a\x08\x03\x12\x08\x12\x02\x08\x01\x12\x02\x08\x03"\x0c\x00\x00\x00\x00\xff\x00\x00\x00\xff\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nD\n\x07Const_3\x12\x05Const*%\n\x05value\x12\x1cB\x1a\x08\x03\x12\x08\x12\x02\x08\x01\x12\x02\x08\x03"\x0c\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nD\n\x07Const_4\x12\x05Const*%\n\x05value\x12\x1cB\x1a\x08\x03\x12\x08\x12\x02\x08\x01\x12\x02\x08\x03"\x0c\x00\x00\x00\x00\xff\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\n5\n\nimg_holder\x12\x0bPlaceholder*\r\n\x05shape\x12\x04:\x02\x18\x01*\x0b\n\x05dtype\x12\x020\x03\n7\n\x0ctarget_color\x12\x0bPlaceholder*\r\n\x05shape\x12\x04:\x02\x18\x01*\x0b\n\x05dtype\x12\x020\x03\nL\n\x13strided_slice/stack\x12\x05Const*!\n\x05value\x12\x18B\x16\x08\x03\x12\x04\x12\x02\x08\x03"\x0c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nN\n\x15strided_slice/stack_1\x12\x05Const*!\n\x05value\x12\x18B\x16\x08\x03\x12\x04\x12\x02\x08\x03"\x0c\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nN\n\x15strided_slice/stack_2\x12\x05Const*!\n\x05value\x12\x18B\x16\x08\x03\x12\x04\x12\x02\x08\x03"\x0c\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\n\xe6\x01\n\rstrided_slice\x12\x0cStridedSlice\x1a\nimg_holder\x1a\x13strided_slice/stack\x1a\x15strided_slice/stack_1\x1a\x15strided_slice/stack_2*\x07\n\x01T\x12\x020\x03*\x0b\n\x05Index\x12\x020\x03*\x16\n\x10shrink_axis_mask\x12\x02\x18\x03*\x10\n\nbegin_mask\x12\x02\x18\x04*\x13\n\rellipsis_mask\x12\x02\x18\x00*\x13\n\rnew_axis_mask\x12\x02\x18\x00*\x0e\n\x08end_mask\x12\x02\x18\x04\nB\n\rReshape/shape\x12\x05Const*\x1d\n\x05value\x12\x14B\x12\x08\x03\x12\x04\x12\x02\x08\x02"\x08\x01\x00\x00\x00\x03\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nG\n\x07Reshape\x12\x07Reshape\x1a\rstrided_slice\x1a\rReshape/shape*\x07\n\x01T\x12\x020\x03*\x0c\n\x06Tshape\x12\x020\x03\nN\n\x15strided_slice_1/stack\x12\x05Const*!\n\x05value\x12\x18B\x16\x08\x03\x12\x04\x12\x02\x08\x03"\x0c\x06\x00\x00\x00\x06\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nP\n\x17strided_slice_1/stack_1\x12\x05Const*!\n\x05value\x12\x18B\x16\x08\x03\x12\x04\x12\x02\x08\x03"\x0c\x07\x00\x00\x00\x07\x00\x00\x00\x00\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nP\n\x17strided_slice_1/stack_2\x12\x05Const*!\n\x05value\x12\x18B\x16\x08\x03\x12\x04\x12\x02\x08\x03"\x0c\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\n\xee\x01\n\x0fstrided_slice_1\x12\x0cStridedSlice\x1a\nimg_holder\x1a\x15strided_slice_1/stack\x1a\x17strided_slice_1/stack_1\x1a\x17strided_slice_1/stack_2*\x07\n\x01T\x12\x020\x03*\x0b\n\x05Index\x12\x020\x03*\x16\n\x10shrink_axis_mask\x12\x02\x18\x03*\x10\n\nbegin_mask\x12\x02\x18\x04*\x13\n\rellipsis_mask\x12\x02\x18\x00*\x13\n\rnew_axis_mask\x12\x02\x18\x00*\x0e\n\x08end_mask\x12\x02\x18\x04\nD\n\x0fReshape_1/shape\x12\x05Const*\x1d\n\x05value\x12\x14B\x12\x08\x03\x12\x04\x12\x02\x08\x02"\x08\x01\x00\x00\x00\x03\x00\x00\x00*\x0b\n\x05dtype\x12\x020\x03\nM\n\tReshape_1\x12\x07Reshape\x1a\x0fstrided_slice_1\x1a\x0fReshape_1/shape*\x07\n\x01T\x12\x020\x03*\x0c\n\x06Tshape\x12\x020\x03\n&\n\x03Sub\x12\x03Sub\x1a\nimg_holder\x1a\x05Const*\x07\n\x01T\x12\x020\x03\n\x18\n\x03Abs\x12\x03Abs\x1a\x03Sub*\x07\n\x01T\x12\x020\x03\n?\n\x15Sum/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nL\n\x03Sum\x12\x03Sum\x1a\x03Abs\x1a\x15Sum/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n9\n\x04Cast\x12\x04Cast\x1a\x03Sum*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n*\n\x05Sub_1\x12\x03Sub\x1a\nimg_holder\x1a\x07Const_1*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_1\x12\x03Abs\x1a\x05Sub_1*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_1/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_1\x12\x03Sum\x1a\x05Abs_1\x1a\x17Sum_1/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_1\x12\x04Cast\x1a\x05Sum_1*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n*\n\x05Sub_2\x12\x03Sub\x1a\nimg_holder\x1a\x07Const_3*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_2\x12\x03Abs\x1a\x05Sub_2*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_2/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_2\x12\x03Sum\x1a\x05Abs_2\x1a\x17Sum_2/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_2\x12\x04Cast\x1a\x05Sum_2*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n*\n\x05Sub_3\x12\x03Sub\x1a\nimg_holder\x1a\x07Const_2*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_3\x12\x03Abs\x1a\x05Sub_3*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_3/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_3\x12\x03Sum\x1a\x05Abs_3\x1a\x17Sum_3/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_3\x12\x04Cast\x1a\x05Sum_3*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n*\n\x05Sub_4\x12\x03Sub\x1a\nimg_holder\x1a\x07Const_2*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_4\x12\x03Abs\x1a\x05Sub_4*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_4/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_4\x12\x03Sum\x1a\x05Abs_4\x1a\x17Sum_4/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_4\x12\x04Cast\x1a\x05Sum_4*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n*\n\x05Sub_5\x12\x03Sub\x1a\nimg_holder\x1a\x07Const_4*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_5\x12\x03Abs\x1a\x05Sub_5*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_5/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_5\x12\x03Sum\x1a\x05Abs_5\x1a\x17Sum_5/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_5\x12\x04Cast\x1a\x05Sum_5*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n*\n\x05Sub_6\x12\x03Sub\x1a\nimg_holder\x1a\x07Reshape*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_6\x12\x03Abs\x1a\x05Sub_6*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_6/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_6\x12\x03Sum\x1a\x05Abs_6\x1a\x17Sum_6/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_6\x12\x04Cast\x1a\x05Sum_6*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n,\n\x05Sub_7\x12\x03Sub\x1a\nimg_holder\x1a\tReshape_1*\x07\n\x01T\x12\x020\x03\n\x1c\n\x05Abs_7\x12\x03Abs\x1a\x05Sub_7*\x07\n\x01T\x12\x020\x03\nA\n\x17Sum_7/reduction_indices\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x02*\x0b\n\x05dtype\x12\x020\x03\nR\n\x05Sum_7\x12\x03Sum\x1a\x05Abs_7\x1a\x17Sum_7/reduction_indices*\n\n\x04Tidx\x12\x020\x03*\x0f\n\tkeep_dims\x12\x02(\x01*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_7\x12\x04Cast\x1a\x05Sum_7*\n\n\x04SrcT\x12\x020\x03*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x01\n>\n\x0bconcat/axis\x12\x05Const*\x1b\n\x05value\x12\x12B\x10\x08\x03\x12\x00:\n\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01*\x0b\n\x05dtype\x12\x020\x03\n{\n\x06concat\x12\x08ConcatV2\x1a\x04Cast\x1a\x06Cast_1\x1a\x06Cast_2\x1a\x06Cast_3\x1a\x06Cast_5\x1a\x06Cast_6\x1a\x06Cast_7\x1a\x06Cast_4\x1a\x0bconcat/axis*\n\n\x04Tidx\x12\x020\x03*\x07\n\x01T\x12\x020\x01*\x07\n\x01N\x12\x02\x18\x08\nC\n\x10ArgMin/dimension\x12\x05Const*\x1b\n\x05value\x12\x12B\x10\x08\x03\x12\x00:\n\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01*\x0b\n\x05dtype\x12\x020\x03\nR\n\x06ArgMin\x12\x06ArgMin\x1a\x06concat\x1a\x10ArgMin/dimension*\n\n\x04Tidx\x12\x020\x03*\x07\n\x01T\x12\x020\x01*\x11\n\x0boutput_type\x12\x020\t\n>\n\x06Cast_8\x12\x04Cast\x1a\x06ArgMin*\n\n\x04SrcT\x12\x020\t*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x03\n-\n\x05Equal\x12\x05Equal\x1a\x06Cast_8\x1a\x0ctarget_color*\x07\n\x01T\x12\x020\x03\n=\n\x06Cast_9\x12\x04Cast\x1a\x05Equal*\n\n\x04SrcT\x12\x020\n*\x0e\n\x08Truncate\x12\x02(\x00*\n\n\x04DstT\x12\x020\x03\n0\n\x05mul/y\x12\x05Const*\x13\n\x05value\x12\nB\x08\x08\x03\x12\x00:\x02\xff\x01*\x0b\n\x05dtype\x12\x020\x03\n"\n\x03mul\x12\x03Mul\x1a\x06Cast_9\x1a\x05mul/y*\x07\n\x01T\x12\x020\x03\n8\n\x0eExpandDims/dim\x12\x05Const*\x12\n\x05value\x12\tB\x07\x08\x03\x12\x00:\x01\x00*\x0b\n\x05dtype\x12\x020\x03\nB\n\nExpandDims\x12\nExpandDims\x1a\x03mul\x1a\x0eExpandDims/dim*\n\n\x04Tdim\x12\x020\x03*\x07\n\x01T\x12\x020\x03\nC\n\x10ExpandDims_1/dim\x12\x05Const*\x1b\n\x05value\x12\x12B\x10\x08\x03\x12\x00:\n\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01*\x0b\n\x05dtype\x12\x020\x03\nM\n\x0cExpandDims_1\x12\nExpandDims\x1a\nExpandDims\x1a\x10ExpandDims_1/dim*\n\n\x04Tdim\x12\x020\x03*\x07\n\x01T\x12\x020\x03\n>\n\x08filtered\x12\x07Squeeze\x1a\x0cExpandDims_1*\x12\n\x0csqueeze_dims\x12\x02\n\x00*\x07\n\x01T\x12\x020\x03\x12\x00'
        self.color_graph = tf.Graph()
        self.color_sess = tf.compat.v1.Session(
            graph=self.color_graph,
            config=tf.compat.v1.ConfigProto(
                allow_soft_placement=True,
                # log_device_placement=True,
                gpu_options=tf.compat.v1.GPUOptions(
                    # allow_growth=True,  # it will cause fragmentation.
                    per_process_gpu_memory_fraction=0.1
                ))
        )
        self.color_graph_def = self.color_graph.as_graph_def()
        self.load_model()
        self.img_holder = self.color_sess.graph.get_tensor_by_name("img_holder:0")
        self.target_color = self.color_sess.graph.get_tensor_by_name("target_color:0")
        self.filtered = self.color_sess.graph.get_tensor_by_name("filtered:0")
        self.color_graph.finalize()

    def load_model(self):
        raw = self.model_raw_v1_14
        self.color_graph_def.ParseFromString(raw)
        with self.color_graph.as_default():
            self.color_sess.run(tf.compat.v1.global_variables_initializer())
            _ = tf.import_graph_def(self.color_graph_def, name="")

    def separate_color(self, image_bytes, color: TargetColor):
        # image = np.asarray(bytearray(image_bytes), dtype="uint8")
        # image = cv2.imdecode(image, -1)
        image = np.array(PilImage.open(io.BytesIO(image_bytes)))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = self.color_sess.run(self.filtered, {self.img_holder: image, self.target_color: color.value})
        mask = bytearray(cv2.imencode('.png', mask)[1])
        return mask


if __name__ == '__main__':

    pass
    # import os
    # source_dir = r'E:\***'
    # target_dir = r'E:\***'
    # if not os.path.exists(target_dir):
    #     os.makedirs(target_dir)
    #
    # source_names = os.listdir(source_dir)
    # color_extract = ColorExtract()
    # st = time.time()
    # for i, name in enumerate(source_names):
    #     img_path = os.path.join(source_dir, name)
    #     if i % 100 == 0:
    #         print(i)
    #     with open(img_path, "rb") as f:
    #         b = f.read()
    #         result = color_extract.separate_color(b, color_map['red'])
    #     target_path = os.path.join(target_dir, name)
    #     with open(target_path, "wb") as f:
    #         f.write(result)
    #
    # print('completed {}'.format(time.time() - st))
